import tensorflow as tf
from PyQt5.QtCore import QThread, pyqtSignal
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
import sys
import cv2
import numpy as np
import os

import TrainThread
import LoadDataThread
import OutputModelThread
import OutputTFLiteModelThread

import MainPage
import AboutPage
# import TutorialPage

class MainWindow(QTabWidget):
    def __init__(self):
        super().__init__()
        self.scale = 0.3
        self.version = ' V1.2.3'
        self.current_status = 'Ready'
        self.isTraining = False
        self.setWindowIcon(QIcon('images/flower.png'))
        self.setWindowTitle('AiSuperTool Free' + self.version)

        self.to_predict_name = ""
        self.class_names = ['类别1', '类别2','类别3','类别4', '类别5']
        self.resize(2800, 1600)
        self.initUI()

        self.train_dir = ""
        self.model = None
        self.train_dataset = None
        self.validate_dataset = None
        self.epochs = 5                # 训练的轮次

        # 创建加载数据集的线程
        self.load_data_thread = LoadDataThread.LoadDataThread(self)
        self.load_data_thread.signal.connect(self.load_data_callback)

        # # 创建训练模型的线程
        self.train_model_thread = TrainThread.TrainThread(self)
        self.train_model_thread.signal.connect(self.train_model_callback)

        # 创建导出模型的线程
        self.output_model_thread = OutputModelThread.OutputModelThread()
        # 创建回调函数
        self.output_model_thread.signal.connect(self.output_model_callback)

        # 创建导出TFLite模型的线程
        self.output_tflite_model_thread = OutputTFLiteModelThread.OutputTFLiteModelThread()
        # 创建回调函数
        self.output_tflite_model_thread.signal.connect(self.output_tflite_model_callback)

    # 窗口关闭事件
    def closeEvent(self, event):
        reply = QMessageBox.question(self,
                                     '退出',
                                     "是否要退出AiSuperTool？",
                                     QMessageBox.Yes | QMessageBox.No,
                                     QMessageBox.No)
        if reply == QMessageBox.Yes:
            self.close()
            event.accept()
        else:
            event.ignore()

    def initUI(self):
        main_widget = MainPage.MainPage(self).create_page()
        about_widget = AboutPage.AboutPage(self).create_page()
        # tutorial_widget = TutorialPage.TutorialPage(self).create_page()

        self.addTab(main_widget, '主页面')
        # self.addTab(tutorial_widget, '教程')
        self.addTab(about_widget, '关于')
        self.setTabIcon(0, QIcon('images/home.png'))
        self.setTabIcon(1, QIcon('images/about.png'))

    # 启动加载数据集的函数
    def start_load(self):
        self.load_data_thread.train_dir = self.train_dir
        self.load_data_thread.start()  # 启动任务线程
        self.label_message.setText('加载数据集中...')
        self.pgb.setValue(0)
        print('加载数据集中...')


    # 加载数据集线程的回调函数
    def load_data_callback(self, number, detail):
        self.label_classes_number.setText('1. 【' + '图片类别数量：' + number + '】')
        self.label_classes_detail.setText('2. 【' + '图片类别详情：' + str(detail) + '】')
        # 更新类别名称
        self.class_names = str(detail)

        fh = open('models/labels.txt', 'w', encoding='utf-8')
        fh.write(str(detail))
        fh.close()

        self.btn_load_dataset.setText("数据集加载完成")
        self.btn_load_dataset.setStyleSheet('''
                                        QPushButton {
                                            width: 550px;
                                            height: 120px;
                                            background-color: gray;
                                            border-radius: 60px;
                                        }
                                        ''')
        self.label_message.setText('数据集加载完成，请点击“开始训练模型”按钮开始训练模型')
        self.pgb.setValue(5)
        print('数据集加载完成，请点击“开始训练模型”按钮开始训练模型')

        self.current_status = 'load_dataset_finish'


    # 启动训练模型的线程
    def start_train(self):
        if (self.model == None):
            self.label_message.setText('未加载数据集，请先加载数据集')
            print('未加载数据集，请先加载数据集')
            return

        if (self.isTraining):
            QMessageBox.about(self, '温馨提示', "训练模型中...")
            print('训练中')
            return

        self.train_model_thread.start()
        self.btn_start_train.setText('训练模型中')
        self.btn_start_train.setStyleSheet('''
                                        QPushButton {
                                            width: 550px;
                                            height: 120px;
                                            background-color: gray;
                                            border-radius: 60px;
                                        }
                                        ''')
        self.label_message.setText('训练模型中...')
        print('训练模型中...')


        self.isTraining = True

    # 训练模型线程回调函数
    def train_model_callback(self, result, loss, accuracy, process):
        self.pgb.setValue(process)
        self.label_loss.setText('3. 【loss 损失值：' + loss + '】')
        self.label_accuracy.setText('4. 【accuracy 准确率：' + accuracy + '】')

        if (result == 'OK'):
            self.label_message.setText('训练模型完成')
            print('训练模型完成')

            # 启动导出模型的线程
            self.output_model()

    # 启动导出模型文件的线程
    def output_model(self):
        self.output_model_thread.model = self.model
        self.output_model_thread.start()
        self.label_message.setText('导出模型中...')

    # 导出模型线程的回调函数
    def output_model_callback(self, result):
        self.label_message.setText('成功导出模型文件！')
        self.pgb.setValue(100)
        print('成功导出模型文件')

        self.isTraining = True

    def output_tflite_model_callback(self, result):
        self.label_message.setText('成功导出模型文件！')
        QMessageBox.about(self, '导出模型文件', "已导出new_model.tflite模型文件到models文件夹下")

        print('成功导出模型文件')

    def load_dataset(self):
        if self.isTraining:
            QMessageBox.about(self, '温馨提示', "训练模型中...")
            print('训练中...')
            return

        # 如果没有加载数据集
        if (self.current_status == 'Ready'):
            self.train_dir = QFileDialog.getExistingDirectory(None, "选取文件夹", "C:/")
            if (self.train_dir == ""):
                print('数据集路径为空')
                return
            # 开启加载数据集的线程
            self.start_load()
        # 如果已经加载了数据集
        if (self.current_status == 'load_dataset_finish'):
            QMessageBox.about(self, '温馨提示', "数据集已经加载完成")

    # 模型推理预测
    def predict_img(self):
        # 判断模型文件是否存在
        if not os.path.exists("models/new_model.h5"):
            QMessageBox.about(self,'模型推理预测', "找不到new_model.h5模型文件，请先训练模型。")
            print('找不到new_model.h5模型文件')
            return
        # 加载模型文件，使用训练出来的模型进行推理预测
        new_model = tf.keras.models.load_model("models/new_model.h5")

        # 判断标签文件是否存在
        if not os.path.exists("models/labels.txt"):
            QMessageBox.about(self, '模型推理预测', "找不到labels.txt标签文件")
            print('找不到标签文件')
            return
        fh = open('models/labels.txt', 'r', encoding='utf-8')
        self.class_names = fh.read().split(',')

        openfile_name = QFileDialog.getOpenFileName(self, '选择要识别的图片', '', 'Image files(*.jpg , *.png)')
        if not os.path.exists(openfile_name[0]):
            print('未选择要识别的图片')
            return

        # 原始图片
        original_image = cv2.imread(openfile_name[0])
        resize_image = cv2.resize(original_image, (224, 224))
        cv2.imwrite("images/show.png", original_image)

        # 显示图片
        self.img_label.setPixmap(QPixmap('images/show.png'))

        img = np.asarray(resize_image)
        outputs = new_model.predict(img.reshape(1, 224, 224, 3))

        # print(outputs)
        result_index = np.argmax(outputs)
        # print(result_index)
        # 索引
        result = self.class_names[result_index].replace('\'', '').replace(' ', '')
        print(result_index)
        print(self.class_names)
        self.result.setText('预测结果：' + '【' + result + '】')
        # 准确率
        accuracy = np.max(outputs)
        # print(accuracy)
        r1 = round(float(accuracy) * 100, 2)
        self.accuracy.setText('准确率：' + format(r1)+'%')


    def output_model_to_computer(self):
        if not os.path.exists("models/new_model.h5"):
            QMessageBox.about(self, '导出模型文件', "找不到new_model.h5模型文件，请先训练模型。")
        else:
            self.label_message.setText('导出模型中...')
            self.output_tflite_model_thread.start()



if __name__ == "__main__":
    app = QApplication(sys.argv)
    x = MainWindow()
    x.show()
    sys.exit(app.exec_())
